import psutil
import os
import sys
import win32file
import wmi
import threading
import queue
import time
import itertools
import tkinter as tk
from tkinter import ttk, messagebox
from collections import deque


class SystemDriveError(Exception):
    def __init__(self, message = "protected system drive"):
        self.message = message
        super().__init__(self.message)


class ZeroDiskError(Exception):
    def __init__(self, message = "physical disk 0"):
        self.message = message
        super().__init__(self.message)


class FixedDiskError(Exception):
    def __init__(self, message = "fixed hard disk media"):
        self.message = message
        super().__init__(self.message)


class Resource:
    def __init__(self, file):
        self.file = file

    @property
    def path(self):
        if hasattr(sys, "_MEIPASS"):
            base = sys._MEIPASS
        else:
            base = os.path.abspath(".")

        return os.path.join(base, self.file)


class ProcScanner:
    def __init__(self, drive):
        self.drive = drive.strip().upper()

        self.path = os.path.abspath(self.drive + ":\\")
        self.processes = []

    def __call__(self, *args, **kwargs):
        return self.scan()

    def scan(self):
        for proc in psutil.process_iter(["pid", "name", "exe", "cmdline"]):
            try:
                if proc.info["exe"] and proc.info["exe"].startswith(self.path):
                    self.processes.append(proc)
                    continue

                if proc.info["cmdline"]:
                    for arg in proc.info["cmdline"]:
                        if arg.startswith(self.path):
                            self.processes.append(proc)
                            break
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                continue

        return self.processes


class DiskAnalysis:
    def __init__(self, drive):
        self.drive = drive.strip().upper()

        self.path = os.path.abspath(self.drive + ":\\")
        self.c = wmi.WMI()

    def __call__(self, *args, **kwargs):
        return self.analyse()

    def analyse(self):
        if self.drive == "C":
            raise SystemDriveError

        got = False
        for disk in self.c.Win32_DiskDrive():
            for partition in disk.associators("Win32_DiskDriveToDiskPartition"):
                for logical in partition.associators("Win32_LogicalDiskToPartition"):
                    if logical.DeviceID[0] == self.drive:
                        got = True
                        break
                if got:
                    break
            if got:
                break

        if disk.Index == 0:
            raise ZeroDiskError

        if (
            win32file.GetDriveType(self.path) != win32file.DRIVE_REMOVABLE
            and disk.MediaType != "External hard disk media"
        ):
            raise FixedDiskError

        return True


class LoadingAnimation:
    def __init__(self, root, size=24, color="#808080"):
        self.root = root
        self.size = size
        self.color = color

        self.frames = ["⣾", "⣽", "⣻", "⢿", "⡿", "⣟", "⣯", "⣷"]
        self.cycle = itertools.cycle(self.frames)
        self.label = tk.Label(
            self.root,
            text="",
            font=("Arial", self.size),
            fg=self.color,
            bg="white",
            borderwidth=0,
            highlightthickness=0
        )
        self.animation_id = None
        self.is_playing = False

    def start(self):
        if not self.is_playing:
            self.is_playing = True
            self.label.place(relx=0.5, rely=0.5, anchor=tk.CENTER)
            self._update_frame()

    def stop(self):
        if self.is_playing:
            if self.animation_id:
                self.label.after_cancel(self.animation_id)
            self.label.place_forget()
            self.is_playing = False

    def _update_frame(self):
        if self.is_playing:
            self.label.config(text=next(self.cycle))
            self.animation_id = self.label.after(100, self._update_frame)

    def set_text(self, text):
        self.label.config(text=text)


class DetailViewer:
    def __init__(self, root, tree):
        self.root = root
        self.tree = tree

        self.ui_interval = 1000
        self.data_interval = 500

        self.windows = {}

        self.data_queue = queue.Queue()
        self.running = True

        self.cpu_history = {}
        self.mem_history = {}

    def _collect(self):
        while self.running:
            for window, (pid, _) in list(self.windows.items()):
                try:
                    proc = psutil.Process(pid)
                    
                    cpu = proc.cpu_percent(interval=0.5)
                    mem = proc.memory_full_info().private / 1024 / 1024
                    
                    if pid not in self.cpu_history:
                        self.cpu_history[pid] = deque(maxlen=5)
                    if pid not in self.mem_history:
                        self.mem_history[pid] = deque(maxlen=3)
                    
                    self.cpu_history[pid].append(cpu)
                    self.mem_history[pid].append(mem)
                    smoothed_cpu = sum(self.cpu_history[pid]) / len(self.cpu_history[pid])
                    smoothed_mem = sum(self.mem_history[pid]) / len(self.mem_history[pid])
                    
                    self.data_queue.put(("update", window, smoothed_cpu, smoothed_mem))
                    
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    self.data_queue.put(("error", window))

                except Exception as ignore:
                    pass
            
            time.sleep(self.data_interval / 1000)

    def _update(self):
        try:
            while not self.data_queue.empty():
                data = self.data_queue.get_nowait()
                
                if data[0] == "update" and data[1] in self.windows:
                    _, window, cpu, mem = data
                    pid, labels = self.windows[window]
                    
                    labels["cpu"].config(text=f"CPU使用率: {cpu:.1f}%")
                    labels["mem"].config(text=f"内存占用: {mem:.2f} MB")
                    labels["status"].config(text=f"最后更新: {time.strftime('%H:%M:%S')}")
                    
                    color = "green"
                    if cpu > 70: color = "red"
                    elif cpu > 30: color = "orange"
                    labels["cpu"].config(foreground=color)
                    
                elif data[0] == "error" and data[1] in self.windows:
                    window = data[1]
                    _, labels = self.windows[window]
                    labels["cpu"].config(text="[进程已终止]")
                    labels["mem"].config(text="[进程已终止]")
                    labels["status"].config(text="状态: 进程已终止", foreground="red")
                    
        except queue.Empty:
            pass
        
        if self.running:
            self.root.after(self.ui_interval, self._update)

    def _stop(self, window):
        if window in self.windows:
            pid, _ = self.windows[window]
            if pid in self.cpu_history:
                del self.cpu_history[pid]
            if pid in self.mem_history:
                del self.mem_history[pid]
            del self.windows[window]
        window.destroy()

    def detail(self, event):
        selected = self.tree.selection()
        if not selected:
            return

        pid, name, path = self.tree.item(selected[0])["values"]

        window = tk.Toplevel(self.root)
        window.title(f"进程详情 - {name}")
        window.iconbitmap(Resource("icon.ico").path)

        labels = {
            "pid": ttk.Label(window, text=f"PID: {pid}"),
            "name": ttk.Label(window, text=f"名称: {name}"),
            "path": ttk.Label(window, text=f"路径: {path}"),
            "cpu": ttk.Label(window, text="CPU使用率: -"),
            "mem": ttk.Label(window, text="内存占用: -"),
            "status": ttk.Label(window, text="状态: 监控中...")
        }
        for label in labels.values():
            label.pack(anchor=tk.W, padx=10, pady=2)

        ttk.Button(window, text="关闭", command=lambda: self._stop(window)).pack(pady=10)
        window.protocol("WM_DELETE_WINDOW", lambda: self._stop(window))

        if not hasattr(self, "windows"):
            self.windows = {}
        self.windows[window] = (pid, labels)

        self._update()

    def start(self):
        self.collect = threading.Thread(target=self._collect, daemon=True)
        self.collect.start()
        self._update()

    def stop(self):
        self.running = False
        if hasattr(self, "collect"):
            self.collect.join(timeout=1)


class ProcessListViewer:
    def __init__(self, root, drive, ignoreDA, scanner):
        self.root = root
        self.drive = drive
        self.ignoreDA = ignoreDA
        self.scanner = scanner

        self.root.geometry("800x600")
        self.root.title(f"进程列表 - {drive} 盘")
        self.root.iconbitmap(Resource("icon.ico").path)
        
        self.processes = []
        self.root.deiconify()
        self.root.protocol("WM_DELETE_WINDOW", self.close)
        
        self.statistic = None
        self.tree = None

        self._info()
        self._interface()
        self._fill()

        self.dv = DetailViewer(self.root, self.tree)
        self.dv.start()
        self._action()

        self.la = LoadingAnimation(self.root)
        self.data_queue = queue.Queue()
        self.root.after(100, self._check_queue)

        self.refresh()

    def _collect(self):
        path = f"{self.drive.upper()}:\\"
        count = sum(1 for proc in self.processes if proc.exe().startswith(path))
        text = (
            f"共 {len(self.processes)} 个进程与 {drive} 盘相关"
            f"（含 {count} 个从该盘启动）"
        )
        return text

    def _fill(self):
        self.tree.delete(*self.tree.get_children())
        
        for proc in self.processes:
            try:
                self.tree.insert(
                    '', 
                    tk.END, 
                    values=(
                        proc.pid, 
                        proc.name(), 
                        proc.exe() or "N/A"
                    )
                )
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue

    def _info(self):
        frame = ttk.Frame(self.root)
        frame.pack(fill=tk.X, padx=10, pady=5)
        
        text = self._collect()
        self.statistic = ttk.Label(frame, text=text, font=("Arial", 10, "bold"))
        self.statistic.pack(side=tk.LEFT)

        if self.ignoreDA:
            ttk.Label(frame, text="\u26a0\ufe0f 不安全模式（ DA 模块已忽略）",
                foreground="red", font=("Arial", 9, "italic")).pack(side=tk.RIGHT)

    def _interface(self):
        self.tree = ttk.Treeview(
            self.root, 
            columns=("PID", "Name", "Path"), 
            show="headings",
            selectmode="browse"
        )

        self.tree.heading("PID", text="PID", anchor=tk.W)
        self.tree.heading("Name", text="进程名称", anchor=tk.W)
        self.tree.heading("Path", text="文件路径", anchor=tk.W)
        
        self.tree.column("PID", width=80, minwidth=50)
        self.tree.column("Name", width=150, minwidth=100)
        self.tree.column("Path", width=500, minwidth=300)

        scrollbar = ttk.Scrollbar(self.root, orient=tk.VERTICAL, command=self.tree.yview)
        self.tree.configure(yscroll=scrollbar.set)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        self.tree.pack(fill=tk.BOTH, expand=True, padx=10, pady=5)

    def _action(self):
        frame = ttk.Frame(self.root)
        frame.pack(fill=tk.X, padx=10, pady=10)
        
        ttk.Button(
            frame, 
            text="终止选中进程", 
            command=self.kill
        ).pack(side=tk.LEFT, padx=5)
        
        self.refresh_btn = ttk.Button(
            frame, 
            text="刷新列表", 
            command=self.refresh
        )
        self.refresh_btn.pack(side=tk.LEFT, padx=5)
        
        ttk.Button(
            frame, 
            text="关闭", 
            command=self.close
        ).pack(side=tk.RIGHT, padx=5)

        self.tree.bind("<Double-1>", self.dv.detail)

    def kill(self):
        selected = self.tree.selection()
        if not selected:
            messagebox.showwarning("警告", "请先选中一个进程")
            return
            
        pid = self.tree.item(selected[0])["values"][0]
        try:
            psutil.Process(pid).kill()
            messagebox.showinfo("成功", f"已终止进程 {pid}")
            self.tree.delete(selected[0])
        except Exception as e:
            messagebox.showerror("错误", f"终止进程失败: {e}")
        finally:
            self.refresh()
    
    def refresh(self):
        self.scanner.processes = []
        self.refresh_btn.config(state=tk.DISABLED)
        self.tree.delete(*self.tree.get_children())
        self.la.start()
        self.statistic.config(text="扫描中...", foreground="blue")

        threading.Thread(
            target=self._async_scanner,
            daemon=True
        ).start()

    def close(self):
        self.dv.stop()
        self.root.destroy()

    def _async_scanner(self):
        try:
            processes = self.scanner()
            self.data_queue.put(("success", processes))
        except Exception as e:
            self.data_queue.put(("error", str(e)))

    def _check_queue(self):
        try:
            while True:
                result = self.data_queue.get_nowait()
                if result[0] == "success":
                    self._handle_success(result[1])
                elif result[0] == "error":
                    self._handle_error(result[1])
        except queue.Empty:
            pass
        
        self.root.after(100, self._check_queue)

    def _handle_success(self, processes):
        self.la.stop()
        
        self.processes = processes
        text = self._collect()
        self.statistic.config(text=text, foreground="black")
        self._fill()
        self.refresh_btn.config(state=tk.NORMAL)

    def _handle_error(self, error_msg):
        self.la.stop()

        self.statistic.config(
            text=f"\u26a0\ufe0f 错误: {error_msg}",
            foreground="red"
        )
        self.refresh_btn.config(state=tk.NORMAL)


class MainApplication:
    def __init__(self, drive, ignoreDA=False):
        self.drive = drive
        self.ignoreDA = ignoreDA

        self.title = "Process Killer"

        self.root = tk.Tk()
        self.root.withdraw()

    def __call__(self, *args, **kwargs):
        return self.main()

    def _initialize(self):
        try:
            if self.ignoreDA is None:
                raise Exception("骗你的，DA 怎么可以跳过，你不要给我喵喵叫")
            _ = DiskAnalysis(drive)()

        except SystemDriveError as e:
            message = (
                f"{e.message.upper()}\n\n"
                "出于系统安全考虑\n"
                "不允许终止从系统盘启动的进程\n"
            )
            icon = messagebox.ERROR
            messagebox.showerror(
                title=self.title, message=message, icon=icon
            )
            sys.exit()

        except ZeroDiskError as e:
            message = (
                f"{e.message.upper()}\n\n"
                f"检测到 {self.drive} 盘属于 磁盘0（系统保留盘）的分区盘\n"
                f"不建议终止从 {self.drive} 盘启动的进程，确定继续吗？\n"
            )
            icon = messagebox.ERROR
            if not messagebox.askokcancel(
                title=self.title, message=message, icon=icon
                , default=messagebox.CANCEL
            ): sys.exit()

        except FixedDiskError as e:
            message = (
                f"{e.message.upper()}\n\n"
                f"检测到 {self.drive} 盘可能是内置硬盘或非标准 USB 设备\n"
                f"没必要终止从 {self.drive} 盘启动的进程，是否继续？\n"
            )
            icon = messagebox.WARNING
            if not messagebox.askyesno(
                title=self.title, message=message, icon=icon
                , default=messagebox.NO
            ): sys.exit()

        except Exception as e:
            message = (
                f"ERROR\n\n{e}\n\n"
                "软件崩溃了喵！发生了什么喵？\n硬盘解析出问题了喵？！\n"
            )
            icon = messagebox.ERROR
            messagebox.showerror(
                title=self.title, message=message, icon=icon
            )
            sys.exit()

    def main(self):
        if self.ignoreDA != True:
            self._initialize()

        scanner = ProcScanner(self.drive)
        plv = ProcessListViewer(
            self.root, self.drive, self.ignoreDA, scanner)
        plv.root.mainloop()


if __name__ == "__main__":
    argv = sys.argv
    if len(argv) < 2:
        if getattr(sys, "frozen", False):
            path = sys.argv[0]
        else:
            path = __file__
        drive = os.path.splitdrive(os.path.abspath(path))[0][0]
    else:
        drive = argv[1].strip().upper()

        path = os.path.abspath(drive + ":\\")
        if not os.path.exists(path):
            sys.exit()

    try:
        MainApplication(drive,
            False if "-da:pass" in argv and "-da:pass:喵喵" in argv else
            None if "-da:pass" in argv else
            True if "-da:pass:喵喵" in argv else
            False
        )()
    except KeyboardInterrupt:
        sys.exit()